{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "01b89c49",
   "metadata": {},
   "source": [
    "# 03 Automatic Differentiation\n",
    "\n",
    "Original Documentation: https://docs.jax.dev/en/latest/automatic-differentiation.html\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c0c35b5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18c78062",
   "metadata": {},
   "source": [
    "## Taking gradients\n",
    "\n",
    "Can differentiate a scalar-valued function with the `jax.grad()` function:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f573e1e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.070650816\n"
     ]
    }
   ],
   "source": [
    "grad_tanh = jax.grad(jnp.tanh)\n",
    "print(grad_tanh(2.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6e55e8b",
   "metadata": {},
   "source": [
    "`jax.grad(f)` takes a function $f$ and returns the gradient of the function, $\\nabla f$.\n",
    "\n",
    "You can chain `jax.grad()` calls to create functions for higher-order derivative.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bae842e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-0.13621868\n",
      "0.25265405\n"
     ]
    }
   ],
   "source": [
    "print(jax.grad(jax.grad(jnp.tanh))(2.0))\n",
    "print(jax.grad(jax.grad(jax.grad(jnp.tanh)))(2.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36032964",
   "metadata": {},
   "source": [
    "Consider the derivatives of the function: $f(x) = x^3 + 2x^2 - 3x + 1$ below\n",
    "\n",
    "- $f^\\prime(x) = 3x^2 + 4x - 3$\n",
    "- $f^{\\prime\\prime}(x) = 6x + 4$\n",
    "- $f^{\\prime\\prime\\prime}(x) = 6$\n",
    "- $f^{\\prime\\prime\\prime\\prime}(x) = 0$\n",
    "\n",
    "We can compute these all in JAX with automatic differentiation:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7761f41e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.0\n",
      "10.0\n",
      "6.0\n",
      "0.0\n"
     ]
    }
   ],
   "source": [
    "f = lambda x: x**3 + 2 * x**2 - 3 * x + 1\n",
    "dfdx = jax.grad(f)\n",
    "d2fdx = jax.grad(dfdx)\n",
    "d3fdx = jax.grad(d2fdx)\n",
    "d4fdx = jax.grad(d3fdx)\n",
    "\n",
    "print(dfdx(1.0))\n",
    "print(d2fdx(1.0))\n",
    "print(d3fdx(1.0))\n",
    "print(d4fdx(1.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00fb3239",
   "metadata": {},
   "source": [
    "## Computing gradients in a linear logistic regression\n",
    "\n",
    "Consider this setup code for a binary logistic regression:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc0eb30d",
   "metadata": {},
   "outputs": [],
   "source": [
    "key = jax.random.key(1701)\n",
    "\n",
    "\n",
    "def sigmoid(x):\n",
    "    \"\"\"\n",
    "    Sigmoid function maps any real-valued input to a probability in (0, 1).\n",
    "    Common use: in binary classification, interpret sigmoid(x) as P(class A).\n",
    "    A typical decision rule: if sigmoid(x) > 0.5, predict class A; else class B.\n",
    "    \"\"\"\n",
    "    return 0.5 * (jnp.tanh(x / 2) + 1)\n",
    "\n",
    "\n",
    "def predict(W, b, inputs):\n",
    "    return sigmoid(jnp.dot(inputs, W) + b)\n",
    "\n",
    "\n",
    "inputs = jnp.array(\n",
    "    [[0.52, 1.12, 0.77], [0.88, -1.08, 0.15], [0.52, 0.06, -1.30], [0.74, -2.49, 1.39]]\n",
    ")\n",
    "targets = jnp.array([True, True, False, True])\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def loss(W, b):\n",
    "    \"\"\"\n",
    "    The binary cross-entropy loss on the predictions with the currents weights and biases.\n",
    "    It is a binary loss since we have two label classes (either True or False).\n",
    "    \"\"\"\n",
    "    preds = predict(W, b, inputs)\n",
    "\n",
    "    # Compute the probability the predictions gave towards the correct label\n",
    "    label_preds = preds * targets + (1 - preds) * (1 - targets)\n",
    "\n",
    "    # Compute the cross-entropy loss between predictions and targets\n",
    "    # cross-entropy loss is defined as sum(-log(prob of correct class))\n",
    "    return jnp.sum(-jnp.log(label_preds))\n",
    "\n",
    "\n",
    "key, W_key, b_key = jax.random.split(key, 3)\n",
    "W = jax.random.normal(key, (3,))\n",
    "b = jax.random.normal(key, ())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d459532",
   "metadata": {},
   "source": [
    "In machine learning, the objective is to minimize the loss. To determine the best direction to adjust the weights, we compute the gradient by differentiating the loss function.\n",
    "\n",
    "Use `jax.grad()` to differentiate `loss()` with respect to its parameters:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f40cd974",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "W_grad=(Array([-1.64406  ,  1.4412102, -1.8369946], dtype=float32),)\n",
      "b_grad=(Array(-2.307187, dtype=float32),)\n",
      "W_grad=Array([-1.64406  ,  1.4412102, -1.8369946], dtype=float32)\n",
      "b_grad=Array(-2.307187, dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# Compute gradient w.r.t W\n",
    "W_grad = jax.grad(loss, argnums=(0,))(W, b)\n",
    "print(f\"{W_grad=}\")\n",
    "\n",
    "# Compute gradient w.r.t b\n",
    "b_grad = jax.grad(loss, argnums=(1,))(W, b)\n",
    "print(f\"{b_grad=}\")\n",
    "\n",
    "# Compute gradient w.r.t W and b\n",
    "W_grad, b_grad = jax.grad(loss, argnums=(0, 1))(W, b)\n",
    "print(f\"{W_grad=}\")\n",
    "print(f\"{b_grad=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b29f5a4",
   "metadata": {},
   "source": [
    "Effectively, `argnums` allows differentiating a function with respect to a set of parameters. For example, `jax.grad(f, i)` is equivalent to $\\frac{\\partial f}{\\partial x_i}$.\n",
    "\n",
    "## **Differentiating with respect to nested lists, tuples, and dicts**\n",
    "\n",
    "`jax.grad()` will also work with dictionaries for parameters. It will return output with respect to each parameter (key) as a separate output dictionary item.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a9c0d610",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'W': Array([-1.64406  ,  1.4412102, -1.8369946], dtype=float32), 'b': Array(-2.307187, dtype=float32)}\n"
     ]
    }
   ],
   "source": [
    "@jax.jit\n",
    "def loss2(params):\n",
    "    preds = predict(params[\"W\"], params[\"b\"], inputs)\n",
    "    label_preds = preds * targets + (1 - preds) * (1 - targets)\n",
    "    return jnp.sum(-jnp.log(label_preds))\n",
    "\n",
    "\n",
    "print(jax.grad(loss2)({\"W\": W, \"b\": b}))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8055c1e",
   "metadata": {},
   "source": [
    "## Evaluating a function and its gradient\n",
    "\n",
    "`jax.value_and_grad()` will compute a function’s value as well as its gradient’s value in one pass.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0b957279",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loss: 5.6301613\n",
      "Confirmation: 5.6301613\n"
     ]
    }
   ],
   "source": [
    "loss_val, Wb_grad = jax.value_and_grad(loss, argnums=(0, 1))(W, b)\n",
    "print(\"Loss:\", loss_val)\n",
    "print(\"Confirmation:\", loss(W, b))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
